#!/usr/bin/env python
# -*- coding: utf-8 -*-

import os
import sys
import time

from htlib import *
from hparser import *
from cluster import ClusterConfig
from barrier import SenseBarrier

class HTestCase:
    def __init__(self, interface, home_dir, cluster_name):
        self.client_no = get_ip_address(interface)
        self.groupid   = -1
        self.switch_cluster(home_dir, cluster_name)

    def switch_cluster(self, home_dir, cluster_name):
        """
            - for scalability test
            - firstly, cancel lock server group if necessary.
        """
        if self.groupid != -1:
            self.barrier.cancel()

        log_out('******switch cluster to %s/%s******\n' % (home_dir, cluster_name))

        self.home_dir     = home_dir
        self.cluster_name = cluster_name

        self.cluster = ClusterConfig(home_dir, cluster_name)
        # print self.cluster

        self.sdir     = self.cluster.conf['cluster_share_dir']
        self.hadoop   = self.cluster.conf['cluster_home_dir']
        self.groupid  = self.cluster.conf['cluster_id']
        self.lockhost = self.cluster.conf['lock_server_host']
        self.lockport = self.cluster.conf['lock_server_port']
        self.host     = self.cluster.conf['ssh_server']
        self.user     = self.cluster.conf['ssh_username']

        self.dnum     = len(self.cluster.get_slaves())
        """ datanode number """
        self.cnum     = len(self.cluster.get_clients())
        """ client number """

        self.lockfile = '%s/TestEnv.lock' % self.sdir
        self.barrier  = SenseBarrier(self.lockhost, self.lockport, self.groupid, self.cnum)
        self.core_parser  = HadoopSiteParser('%s/conf/core-site.xml' % self.cluster.conf['cluster_home_dir'])
        self.hdfs_parser  = HadoopSiteParser('%s/conf/hdfs-site.xml' % self.cluster.conf['cluster_home_dir'])

        nfs_mkdirs(self.sdir)

    def finish_barrier(self):
        log_out('exit from group %d\n' % self.groupid)
        self.barrier.cancel()

    def lock(self):
        """ 
        Open with os.O_EXCL is a atomic operation.

        NOTE: In a lock block, can not call barrier, barrier must be called by all clients.
        """
        had_lock = True
        try:
            fd = os.open(self.lockfile, os.O_CREAT|os.O_EXCL)
            os.close(fd)
        except OSError, e:
            had_lock = False

        return had_lock

    def unlock(self):
        try:
            os.unlink(self.lockfile)
        except OSError, e:
            pass 

    def prepare(self):
        """
        virtual function
        """
        pass

    def finish(self):
        """
        virtual function
        """
        pass

    def run(self):
        """
        virtual function
        """
        pass

    def start(self):
        """
        template method
        """
        try:
            self.prepare()
            self.run()
            self.finish()
            self.finish_barrier()
        except KeyboardInterrupt:
            self.finish_barrier()

class HTestEnv:
    def __init__(self, home_dir, cluster_name):
        self.home_dir      = home_dir
        self.cluster_name  = cluster_name

        self.cluster = ClusterConfig(home_dir, cluster_name)

        self.ssh_server      = self.cluster.conf['ssh_server']
        self.user          = self.cluster.conf['ssh_username']
        self.reboot_script = '%s/hadoop-test/script/reboot_hdp.py' % home_dir

    def set_sshserver(self, address):
        self.ssh_server = address

    def copy_conf(self, source, dist):
        command = 'cp -f %s %s' % (source, dist)
        log_out('%s\n' % command)
        try:
            os.system(command)
        except:
            raise

    def stop_hdp(self):
        command = "ssh %s@%s '%s -H %s -c %s'" % (self.user, self.ssh_server, self.reboot_script, self.home_dir, self.cluster_name)
        os.system(command)

    def restart_hdp(self):
        command = "ssh %s@%s '%s -H %s -c %s -r'" % (self.user, self.ssh_server, self.reboot_script, self.home_dir, self.cluster_name)
        os.system(command)

    def restart_jobtracker(self):
        command = "ssh %s@%s '%s -H %s -c %s -j'" % (self.user, self.ssh_server, self.reboot_script, self.home_dir, self.cluster_name)
        os.system(command)

    def format_hdp(self, restart=True, all=False):
        """
        @param restart: restart cluster or not
        """
        if restart:
            command = "ssh %s@%s '%s -H %s -c %s -f -r'" % (self.user, self.ssh_server, self.reboot_script, self.home_dir, self.cluster_name)
        else:
            command = "ssh %s@%s '%s -H %s -c %s -f'" % (self.user, self.ssh_server, self.reboot_script, self.home_dir, self.cluster_name)

        if all:
            command += ' -a'

        os.system(command)

